import json

from selenium import webdriver
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities

#http://www.sahajamit.com/post/capturing-websocket-messages-using-selenium/
#https://stackoverflow.com/questions/27644615/getting-chrome-performance-and-tracing-logs


def create_driver():
  caps = DesiredCapabilities.CHROME
  caps['loggingPrefs'] = {'performance': 'ALL'}
  driver = webdriver.Chrome('bin/chromedriver', desired_capabilities=caps)
  driver.implicitly_wait(3)
  driver.get('https://upbit.com')
  driver.implicitly_wait(8)
  driver.get("https://upbit.com/exchange?code=CRIX.UPBIT.KRW-BTC")
  return (driver)


def strip_left(line, prefix):
  if line.startswith(prefix):
    return line[len(prefix):]
  else:
    return line


def gen_websocket_received(driver):
  for entry in driver.get_log("performance"):
    message_str = entry['message']
    message = json.loads(message_str)
    timestamp = message['message']["params"]["timestamp"]
    method = message['message']["method"]
    if method.lower() == "Network.webSocketFrameReceived".lower():
      payload_str = message['message']['params']['response']['payloadData']
      if payload_str.startswith("a"):
        payload_str = strip_left(payload_str, "a")
        payload = json.loads(payload_str)
        for payload_elem_str in payload:
          payload_elem = json.loads(payload_elem_str)
          yield payload_elem
      elif payload_str.startswith("h"):
        pass
      else:
        pass


driver = create_driver()

while True:
  for entry in gen_websocket_received(driver):
    type = entry["type"]
    if type == "crixOrderbook":
      print("book", entry)
    elif type == "crixTrade":
      print("trade", entry)
    else:
      pass
